{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = ''\n",
    "os.environ['GOOGLE_APPLICATION_CREDENTIALS'] = 'prepare/mesolitica-tpu.json'\n",
    "b2_application_key_id = os.environ['b2_application_key_id']\n",
    "b2_application_key = os.environ['b2_application_key']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from google.cloud import storage\n",
    "client = storage.Client()\n",
    "bucket = client.bucket('mesolitica-tpu-general')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "best = '1100000'\n",
    "directory = 't5-small-spelling'\n",
    "!rm -rf output out {directory}\n",
    "!mkdir {directory}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = best\n",
    "\n",
    "blob = bucket.blob(f'{directory}/model.ckpt-{model}.data-00000-of-00002')\n",
    "blob.download_to_filename(f'{directory}/model.ckpt-{model}.data-00000-of-00002')\n",
    "\n",
    "blob = bucket.blob(f'{directory}/model.ckpt-{model}.data-00001-of-00002')\n",
    "blob.download_to_filename(f'{directory}/model.ckpt-{model}.data-00001-of-00002')\n",
    "\n",
    "blob = bucket.blob(f'{directory}/model.ckpt-{model}.index')\n",
    "blob.download_to_filename(f'{directory}/model.ckpt-{model}.index')\n",
    "\n",
    "blob = bucket.blob(f'{directory}/model.ckpt-{model}.meta')\n",
    "blob.download_to_filename(f'{directory}/model.ckpt-{model}.meta')\n",
    "\n",
    "blob = bucket.blob(f'{directory}/checkpoint')\n",
    "blob.download_to_filename(f'{directory}/checkpoint')\n",
    "\n",
    "blob = bucket.blob(f'{directory}/operative_config.gin')\n",
    "blob.download_to_filename(f'{directory}/operative_config.gin')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(f'{directory}/checkpoint', 'w') as fopen:\n",
    "    fopen.write(f'model_checkpoint_path: \"model.ckpt-{model}\"')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "from b2sdk.v1 import *\n",
    "info = InMemoryAccountInfo()\n",
    "b2_api = B2Api(info)\n",
    "application_key_id = b2_application_key_id\n",
    "application_key = b2_application_key\n",
    "b2_api.authorize_account(\"production\", application_key_id, application_key)\n",
    "file_info = {'how': 'good-file'}\n",
    "b2_bucket = b2_api.get_bucket_by_name('malaya-model')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<b2sdk.file_version.FileVersionInfo at 0x7f7a4688aba8>"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tar = 't5-small-spelling-correction-2021-09-10.tar.gz'\n",
    "os.system(f'tar -czvf {tar} {directory}')\n",
    "outPutname = f'finetuned/{tar}'\n",
    "b2_bucket.upload_local_file(\n",
    "    local_file=tar,\n",
    "    file_name=outPutname,\n",
    "    file_infos=file_info,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "os.system(f'rm {tar}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import tensorflow_datasets as tfds\n",
    "import t5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = t5.models.MtfModel(\n",
    "    model_dir=directory,\n",
    "    tpu=None,\n",
    "    tpu_topology=None,\n",
    "    model_parallelism=1,\n",
    "    batch_size=1,\n",
    "    sequence_length={\"inputs\": 256, \"targets\": 256},\n",
    "    learning_rate_schedule=0.003,\n",
    "    save_checkpoints_steps=5000,\n",
    "    keep_checkpoint_max=3,\n",
    "    iterations_per_loop=100,\n",
    "    mesh_shape=\"model:1,batch:1\", \n",
    "    mesh_devices=[\"cpu:0\"]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "!rm -rf output/*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Using config: {'_model_dir': 't5-small-spelling', '_tf_random_seed': None, '_save_summary_steps': 100, '_save_checkpoints_steps': None, '_save_checkpoints_secs': None, '_session_config': allow_soft_placement: true\n",
      "isolate_session_state: true\n",
      ", '_keep_checkpoint_max': 5, '_keep_checkpoint_every_n_hours': 10000, '_log_step_count_steps': None, '_train_distribute': None, '_device_fn': None, '_protocol': None, '_eval_distribute': None, '_experimental_distribute': None, '_experimental_max_worker_delay_secs': None, '_session_creation_timeout_secs': 7200, '_service': None, '_cluster_spec': None, '_task_type': 'worker', '_task_id': 0, '_global_id_in_cluster': 0, '_master': '', '_evaluation_master': '', '_is_chief': True, '_num_ps_replicas': 0, '_num_worker_replicas': 1, '_tpu_config': TPUConfig(iterations_per_loop=100, num_shards=None, num_cores_per_replica=1, per_host_input_for_training=4, tpu_job_name=None, initial_infeed_sleep_secs=None, input_partition_dims=None, eval_training_input_configuration=2, experimental_host_call_every_n_steps=1), '_cluster': <tensorflow.python.distribute.cluster_resolver.tpu_cluster_resolver.TPUClusterResolver object at 0x7f793b05d748>}\n",
      "INFO:tensorflow:_TPUContext: eval_on_tpu True\n",
      "WARNING:tensorflow:eval_on_tpu ignored because use_tpu is False.\n"
     ]
    }
   ],
   "source": [
    "import gin\n",
    "\n",
    "from t5.data import sentencepiece_vocabulary\n",
    "\n",
    "DEFAULT_SPM_PATH = 'prepare/sp10m.cased.ms-en.model'\n",
    "DEFAULT_EXTRA_IDS = 100\n",
    "model_dir = directory\n",
    "\n",
    "def get_default_vocabulary():\n",
    "    return sentencepiece_vocabulary.SentencePieceVocabulary(\n",
    "      DEFAULT_SPM_PATH, DEFAULT_EXTRA_IDS)\n",
    "\n",
    "with gin.unlock_config():\n",
    "    gin.parse_config_file(t5.models.mtf_model._operative_config_path(model_dir))\n",
    "    gin.bind_parameter(\"Bitransformer.decode.beam_size\", 1)\n",
    "    gin.bind_parameter(\"Bitransformer.decode.temperature\", 0)\n",
    "    gin.bind_parameter(\"utils.get_variable_dtype.slice_dtype\", \"float32\")\n",
    "    gin.bind_parameter(\n",
    "        \"utils.get_variable_dtype.activation_dtype\", \"float32\")\n",
    "    \n",
    "vocabulary = t5.data.SentencePieceVocabulary(DEFAULT_SPM_PATH)\n",
    "estimator = model.estimator(vocabulary, disable_tpu=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(1100000, 'model.ckpt-1100000', 't5-small-spelling/model.ckpt-1100000')"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import os\n",
    "\n",
    "checkpoint_step = t5.models.mtf_model._get_latest_checkpoint_from_dir(model_dir)\n",
    "model_ckpt = \"model.ckpt-\" + str(checkpoint_step)\n",
    "checkpoint_path = os.path.join(model_dir, model_ckpt)\n",
    "checkpoint_step, model_ckpt, checkpoint_path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /home/husein/.local/lib/python3.6/site-packages/tensorflow_core/python/ops/resource_variable_ops.py:1630: calling BaseResourceVariable.__init__ (from tensorflow.python.ops.resource_variable_ops) with constraint is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "If using Keras pass *_constraint arguments to layers.\n",
      "INFO:tensorflow:Calling model_fn.\n",
      "INFO:tensorflow:Running infer on CPU\n",
      "INFO:tensorflow:feature inputs : Tensor(\"Reshape:0\", shape=(1, 256), dtype=int32)\n",
      "WARNING:tensorflow:From /home/husein/.local/lib/python3.6/site-packages/mesh_tensorflow/transformer/utils.py:427: Print (from tensorflow.python.ops.logging_ops) is deprecated and will be removed after 2018-08-20.\n",
      "Instructions for updating:\n",
      "Use tf.print instead of tf.Print. Note that tf.print returns a no-output operator that directly prints the output. Outside of defuns or eager mode, this operator will not be executed unless it is directly specified in session.run or used as a control dependency for other operators. This is only a concern in graph mode. Below is an example of how to ensure tf.print executes in graph mode:\n",
      "\n",
      "WARNING:tensorflow:Using default tf glorot_uniform_initializer for variable encoder/block_000/layer_000/SelfAttention/relative_attention_bias  The initialzer will guess the input and output dimensions  based on dimension order.\n",
      "WARNING:tensorflow:Using default tf glorot_uniform_initializer for variable decoder/block_000/layer_000/SelfAttention/relative_attention_bias  The initialzer will guess the input and output dimensions  based on dimension order.\n",
      "WARNING:tensorflow:Using default tf glorot_uniform_initializer for variable decoder/block_000/layer_000/SelfAttention/relative_attention_bias  The initialzer will guess the input and output dimensions  based on dimension order.\n",
      "INFO:tensorflow:Variable decoder/block_000/layer_000/SelfAttention/k                  size 262144       slice_size 262144       Shape[d_model=512, heads=512]                               \n",
      "INFO:tensorflow:Variable decoder/block_000/layer_000/SelfAttention/o                  size 262144       slice_size 262144       Shape[heads=512, d_model=512]                               \n",
      "INFO:tensorflow:Variable decoder/block_000/layer_000/SelfAttention/q                  size 262144       slice_size 262144       Shape[d_model=512, heads=512]                               \n",
      "INFO:tensorflow:Variable decoder/block_000/layer_000/SelfAttention/relative_attention_bias size 256          slice_size 256          Shape[heads=8, buckets=32]                                  \n",
      "INFO:tensorflow:Variable decoder/block_000/layer_000/SelfAttention/v                  size 262144       slice_size 262144       Shape[d_model=512, heads=512]                               \n",
      "INFO:tensorflow:Variable decoder/block_000/layer_000/layer_norm/scale                 size 512          slice_size 512          Shape[d_model=512]                                          \n",
      "INFO:tensorflow:Variable decoder/block_000/layer_001/EncDecAttention/k                size 262144       slice_size 262144       Shape[d_model=512, heads=512]                               \n",
      "INFO:tensorflow:Variable decoder/block_000/layer_001/EncDecAttention/o                size 262144       slice_size 262144       Shape[heads=512, d_model=512]                               \n",
      "INFO:tensorflow:Variable decoder/block_000/layer_001/EncDecAttention/q                size 262144       slice_size 262144       Shape[d_model=512, heads=512]                               \n",
      "INFO:tensorflow:Variable decoder/block_000/layer_001/EncDecAttention/v                size 262144       slice_size 262144       Shape[d_model=512, heads=512]                               \n",
      "INFO:tensorflow:Variable decoder/block_000/layer_001/layer_norm/scale                 size 512          slice_size 512          Shape[d_model=512]                                          \n",
      "INFO:tensorflow:Variable decoder/block_000/layer_002/DenseReluDense/wi/kernel         size 1048576      slice_size 1048576      Shape[d_model=512, d_ff=2048]                               \n",
      "INFO:tensorflow:Variable decoder/block_000/layer_002/DenseReluDense/wo/kernel         size 1048576      slice_size 1048576      Shape[d_ff=2048, d_model=512]                               \n",
      "INFO:tensorflow:Variable decoder/block_000/layer_002/layer_norm/scale                 size 512          slice_size 512          Shape[d_model=512]                                          \n",
      "INFO:tensorflow:Variable decoder/block_001/layer_000/SelfAttention/k                  size 262144       slice_size 262144       Shape[d_model=512, heads=512]                               \n",
      "INFO:tensorflow:Variable decoder/block_001/layer_000/SelfAttention/o                  size 262144       slice_size 262144       Shape[heads=512, d_model=512]                               \n",
      "INFO:tensorflow:Variable decoder/block_001/layer_000/SelfAttention/q                  size 262144       slice_size 262144       Shape[d_model=512, heads=512]                               \n",
      "INFO:tensorflow:Variable decoder/block_001/layer_000/SelfAttention/v                  size 262144       slice_size 262144       Shape[d_model=512, heads=512]                               \n",
      "INFO:tensorflow:Variable decoder/block_001/layer_000/layer_norm/scale                 size 512          slice_size 512          Shape[d_model=512]                                          \n",
      "INFO:tensorflow:Variable decoder/block_001/layer_001/EncDecAttention/k                size 262144       slice_size 262144       Shape[d_model=512, heads=512]                               \n",
      "INFO:tensorflow:Variable decoder/block_001/layer_001/EncDecAttention/o                size 262144       slice_size 262144       Shape[heads=512, d_model=512]                               \n",
      "INFO:tensorflow:Variable decoder/block_001/layer_001/EncDecAttention/q                size 262144       slice_size 262144       Shape[d_model=512, heads=512]                               \n",
      "INFO:tensorflow:Variable decoder/block_001/layer_001/EncDecAttention/v                size 262144       slice_size 262144       Shape[d_model=512, heads=512]                               \n",
      "INFO:tensorflow:Variable decoder/block_001/layer_001/layer_norm/scale                 size 512          slice_size 512          Shape[d_model=512]                                          \n",
      "INFO:tensorflow:Variable decoder/block_001/layer_002/DenseReluDense/wi/kernel         size 1048576      slice_size 1048576      Shape[d_model=512, d_ff=2048]                               \n",
      "INFO:tensorflow:Variable decoder/block_001/layer_002/DenseReluDense/wo/kernel         size 1048576      slice_size 1048576      Shape[d_ff=2048, d_model=512]                               \n",
      "INFO:tensorflow:Variable decoder/block_001/layer_002/layer_norm/scale                 size 512          slice_size 512          Shape[d_model=512]                                          \n",
      "INFO:tensorflow:Variable decoder/block_002/layer_000/SelfAttention/k                  size 262144       slice_size 262144       Shape[d_model=512, heads=512]                               \n",
      "INFO:tensorflow:Variable decoder/block_002/layer_000/SelfAttention/o                  size 262144       slice_size 262144       Shape[heads=512, d_model=512]                               \n",
      "INFO:tensorflow:Variable decoder/block_002/layer_000/SelfAttention/q                  size 262144       slice_size 262144       Shape[d_model=512, heads=512]                               \n",
      "INFO:tensorflow:Variable decoder/block_002/layer_000/SelfAttention/v                  size 262144       slice_size 262144       Shape[d_model=512, heads=512]                               \n",
      "INFO:tensorflow:Variable decoder/block_002/layer_000/layer_norm/scale                 size 512          slice_size 512          Shape[d_model=512]                                          \n",
      "INFO:tensorflow:Variable decoder/block_002/layer_001/EncDecAttention/k                size 262144       slice_size 262144       Shape[d_model=512, heads=512]                               \n",
      "INFO:tensorflow:Variable decoder/block_002/layer_001/EncDecAttention/o                size 262144       slice_size 262144       Shape[heads=512, d_model=512]                               \n",
      "INFO:tensorflow:Variable decoder/block_002/layer_001/EncDecAttention/q                size 262144       slice_size 262144       Shape[d_model=512, heads=512]                               \n",
      "INFO:tensorflow:Variable decoder/block_002/layer_001/EncDecAttention/v                size 262144       slice_size 262144       Shape[d_model=512, heads=512]                               \n",
      "INFO:tensorflow:Variable decoder/block_002/layer_001/layer_norm/scale                 size 512          slice_size 512          Shape[d_model=512]                                          \n",
      "INFO:tensorflow:Variable decoder/block_002/layer_002/DenseReluDense/wi/kernel         size 1048576      slice_size 1048576      Shape[d_model=512, d_ff=2048]                               \n",
      "INFO:tensorflow:Variable decoder/block_002/layer_002/DenseReluDense/wo/kernel         size 1048576      slice_size 1048576      Shape[d_ff=2048, d_model=512]                               \n",
      "INFO:tensorflow:Variable decoder/block_002/layer_002/layer_norm/scale                 size 512          slice_size 512          Shape[d_model=512]                                          \n",
      "INFO:tensorflow:Variable decoder/block_003/layer_000/SelfAttention/k                  size 262144       slice_size 262144       Shape[d_model=512, heads=512]                               \n",
      "INFO:tensorflow:Variable decoder/block_003/layer_000/SelfAttention/o                  size 262144       slice_size 262144       Shape[heads=512, d_model=512]                               \n",
      "INFO:tensorflow:Variable decoder/block_003/layer_000/SelfAttention/q                  size 262144       slice_size 262144       Shape[d_model=512, heads=512]                               \n",
      "INFO:tensorflow:Variable decoder/block_003/layer_000/SelfAttention/v                  size 262144       slice_size 262144       Shape[d_model=512, heads=512]                               \n",
      "INFO:tensorflow:Variable decoder/block_003/layer_000/layer_norm/scale                 size 512          slice_size 512          Shape[d_model=512]                                          \n",
      "INFO:tensorflow:Variable decoder/block_003/layer_001/EncDecAttention/k                size 262144       slice_size 262144       Shape[d_model=512, heads=512]                               \n",
      "INFO:tensorflow:Variable decoder/block_003/layer_001/EncDecAttention/o                size 262144       slice_size 262144       Shape[heads=512, d_model=512]                               \n",
      "INFO:tensorflow:Variable decoder/block_003/layer_001/EncDecAttention/q                size 262144       slice_size 262144       Shape[d_model=512, heads=512]                               \n",
      "INFO:tensorflow:Variable decoder/block_003/layer_001/EncDecAttention/v                size 262144       slice_size 262144       Shape[d_model=512, heads=512]                               \n",
      "INFO:tensorflow:Variable decoder/block_003/layer_001/layer_norm/scale                 size 512          slice_size 512          Shape[d_model=512]                                          \n",
      "INFO:tensorflow:Variable decoder/block_003/layer_002/DenseReluDense/wi/kernel         size 1048576      slice_size 1048576      Shape[d_model=512, d_ff=2048]                               \n",
      "INFO:tensorflow:Variable decoder/block_003/layer_002/DenseReluDense/wo/kernel         size 1048576      slice_size 1048576      Shape[d_ff=2048, d_model=512]                               \n",
      "INFO:tensorflow:Variable decoder/block_003/layer_002/layer_norm/scale                 size 512          slice_size 512          Shape[d_model=512]                                          \n",
      "INFO:tensorflow:Variable decoder/block_004/layer_000/SelfAttention/k                  size 262144       slice_size 262144       Shape[d_model=512, heads=512]                               \n",
      "INFO:tensorflow:Variable decoder/block_004/layer_000/SelfAttention/o                  size 262144       slice_size 262144       Shape[heads=512, d_model=512]                               \n",
      "INFO:tensorflow:Variable decoder/block_004/layer_000/SelfAttention/q                  size 262144       slice_size 262144       Shape[d_model=512, heads=512]                               \n",
      "INFO:tensorflow:Variable decoder/block_004/layer_000/SelfAttention/v                  size 262144       slice_size 262144       Shape[d_model=512, heads=512]                               \n",
      "INFO:tensorflow:Variable decoder/block_004/layer_000/layer_norm/scale                 size 512          slice_size 512          Shape[d_model=512]                                          \n",
      "INFO:tensorflow:Variable decoder/block_004/layer_001/EncDecAttention/k                size 262144       slice_size 262144       Shape[d_model=512, heads=512]                               \n",
      "INFO:tensorflow:Variable decoder/block_004/layer_001/EncDecAttention/o                size 262144       slice_size 262144       Shape[heads=512, d_model=512]                               \n",
      "INFO:tensorflow:Variable decoder/block_004/layer_001/EncDecAttention/q                size 262144       slice_size 262144       Shape[d_model=512, heads=512]                               \n",
      "INFO:tensorflow:Variable decoder/block_004/layer_001/EncDecAttention/v                size 262144       slice_size 262144       Shape[d_model=512, heads=512]                               \n",
      "INFO:tensorflow:Variable decoder/block_004/layer_001/layer_norm/scale                 size 512          slice_size 512          Shape[d_model=512]                                          \n",
      "INFO:tensorflow:Variable decoder/block_004/layer_002/DenseReluDense/wi/kernel         size 1048576      slice_size 1048576      Shape[d_model=512, d_ff=2048]                               \n",
      "INFO:tensorflow:Variable decoder/block_004/layer_002/DenseReluDense/wo/kernel         size 1048576      slice_size 1048576      Shape[d_ff=2048, d_model=512]                               \n",
      "INFO:tensorflow:Variable decoder/block_004/layer_002/layer_norm/scale                 size 512          slice_size 512          Shape[d_model=512]                                          \n",
      "INFO:tensorflow:Variable decoder/block_005/layer_000/SelfAttention/k                  size 262144       slice_size 262144       Shape[d_model=512, heads=512]                               \n",
      "INFO:tensorflow:Variable decoder/block_005/layer_000/SelfAttention/o                  size 262144       slice_size 262144       Shape[heads=512, d_model=512]                               \n",
      "INFO:tensorflow:Variable decoder/block_005/layer_000/SelfAttention/q                  size 262144       slice_size 262144       Shape[d_model=512, heads=512]                               \n",
      "INFO:tensorflow:Variable decoder/block_005/layer_000/SelfAttention/v                  size 262144       slice_size 262144       Shape[d_model=512, heads=512]                               \n",
      "INFO:tensorflow:Variable decoder/block_005/layer_000/layer_norm/scale                 size 512          slice_size 512          Shape[d_model=512]                                          \n",
      "INFO:tensorflow:Variable decoder/block_005/layer_001/EncDecAttention/k                size 262144       slice_size 262144       Shape[d_model=512, heads=512]                               \n",
      "INFO:tensorflow:Variable decoder/block_005/layer_001/EncDecAttention/o                size 262144       slice_size 262144       Shape[heads=512, d_model=512]                               \n",
      "INFO:tensorflow:Variable decoder/block_005/layer_001/EncDecAttention/q                size 262144       slice_size 262144       Shape[d_model=512, heads=512]                               \n",
      "INFO:tensorflow:Variable decoder/block_005/layer_001/EncDecAttention/v                size 262144       slice_size 262144       Shape[d_model=512, heads=512]                               \n",
      "INFO:tensorflow:Variable decoder/block_005/layer_001/layer_norm/scale                 size 512          slice_size 512          Shape[d_model=512]                                          \n",
      "INFO:tensorflow:Variable decoder/block_005/layer_002/DenseReluDense/wi/kernel         size 1048576      slice_size 1048576      Shape[d_model=512, d_ff=2048]                               \n",
      "INFO:tensorflow:Variable decoder/block_005/layer_002/DenseReluDense/wo/kernel         size 1048576      slice_size 1048576      Shape[d_ff=2048, d_model=512]                               \n",
      "INFO:tensorflow:Variable decoder/block_005/layer_002/layer_norm/scale                 size 512          slice_size 512          Shape[d_model=512]                                          \n",
      "INFO:tensorflow:Variable decoder/final_layer_norm/scale                               size 512          slice_size 512          Shape[d_model=512]                                          \n",
      "INFO:tensorflow:Variable encoder/block_000/layer_000/SelfAttention/k                  size 262144       slice_size 262144       Shape[d_model=512, heads=512]                               \n",
      "INFO:tensorflow:Variable encoder/block_000/layer_000/SelfAttention/o                  size 262144       slice_size 262144       Shape[heads=512, d_model=512]                               \n",
      "INFO:tensorflow:Variable encoder/block_000/layer_000/SelfAttention/q                  size 262144       slice_size 262144       Shape[d_model=512, heads=512]                               \n",
      "INFO:tensorflow:Variable encoder/block_000/layer_000/SelfAttention/relative_attention_bias size 256          slice_size 256          Shape[heads=8, buckets=32]                                  \n",
      "INFO:tensorflow:Variable encoder/block_000/layer_000/SelfAttention/v                  size 262144       slice_size 262144       Shape[d_model=512, heads=512]                               \n",
      "INFO:tensorflow:Variable encoder/block_000/layer_000/layer_norm/scale                 size 512          slice_size 512          Shape[d_model=512]                                          \n",
      "INFO:tensorflow:Variable encoder/block_000/layer_001/DenseReluDense/wi/kernel         size 1048576      slice_size 1048576      Shape[d_model=512, d_ff=2048]                               \n",
      "INFO:tensorflow:Variable encoder/block_000/layer_001/DenseReluDense/wo/kernel         size 1048576      slice_size 1048576      Shape[d_ff=2048, d_model=512]                               \n",
      "INFO:tensorflow:Variable encoder/block_000/layer_001/layer_norm/scale                 size 512          slice_size 512          Shape[d_model=512]                                          \n",
      "INFO:tensorflow:Variable encoder/block_001/layer_000/SelfAttention/k                  size 262144       slice_size 262144       Shape[d_model=512, heads=512]                               \n",
      "INFO:tensorflow:Variable encoder/block_001/layer_000/SelfAttention/o                  size 262144       slice_size 262144       Shape[heads=512, d_model=512]                               \n",
      "INFO:tensorflow:Variable encoder/block_001/layer_000/SelfAttention/q                  size 262144       slice_size 262144       Shape[d_model=512, heads=512]                               \n",
      "INFO:tensorflow:Variable encoder/block_001/layer_000/SelfAttention/v                  size 262144       slice_size 262144       Shape[d_model=512, heads=512]                               \n",
      "INFO:tensorflow:Variable encoder/block_001/layer_000/layer_norm/scale                 size 512          slice_size 512          Shape[d_model=512]                                          \n",
      "INFO:tensorflow:Variable encoder/block_001/layer_001/DenseReluDense/wi/kernel         size 1048576      slice_size 1048576      Shape[d_model=512, d_ff=2048]                               \n",
      "INFO:tensorflow:Variable encoder/block_001/layer_001/DenseReluDense/wo/kernel         size 1048576      slice_size 1048576      Shape[d_ff=2048, d_model=512]                               \n",
      "INFO:tensorflow:Variable encoder/block_001/layer_001/layer_norm/scale                 size 512          slice_size 512          Shape[d_model=512]                                          \n",
      "INFO:tensorflow:Variable encoder/block_002/layer_000/SelfAttention/k                  size 262144       slice_size 262144       Shape[d_model=512, heads=512]                               \n",
      "INFO:tensorflow:Variable encoder/block_002/layer_000/SelfAttention/o                  size 262144       slice_size 262144       Shape[heads=512, d_model=512]                               \n",
      "INFO:tensorflow:Variable encoder/block_002/layer_000/SelfAttention/q                  size 262144       slice_size 262144       Shape[d_model=512, heads=512]                               \n",
      "INFO:tensorflow:Variable encoder/block_002/layer_000/SelfAttention/v                  size 262144       slice_size 262144       Shape[d_model=512, heads=512]                               \n",
      "INFO:tensorflow:Variable encoder/block_002/layer_000/layer_norm/scale                 size 512          slice_size 512          Shape[d_model=512]                                          \n",
      "INFO:tensorflow:Variable encoder/block_002/layer_001/DenseReluDense/wi/kernel         size 1048576      slice_size 1048576      Shape[d_model=512, d_ff=2048]                               \n",
      "INFO:tensorflow:Variable encoder/block_002/layer_001/DenseReluDense/wo/kernel         size 1048576      slice_size 1048576      Shape[d_ff=2048, d_model=512]                               \n",
      "INFO:tensorflow:Variable encoder/block_002/layer_001/layer_norm/scale                 size 512          slice_size 512          Shape[d_model=512]                                          \n",
      "INFO:tensorflow:Variable encoder/block_003/layer_000/SelfAttention/k                  size 262144       slice_size 262144       Shape[d_model=512, heads=512]                               \n",
      "INFO:tensorflow:Variable encoder/block_003/layer_000/SelfAttention/o                  size 262144       slice_size 262144       Shape[heads=512, d_model=512]                               \n",
      "INFO:tensorflow:Variable encoder/block_003/layer_000/SelfAttention/q                  size 262144       slice_size 262144       Shape[d_model=512, heads=512]                               \n",
      "INFO:tensorflow:Variable encoder/block_003/layer_000/SelfAttention/v                  size 262144       slice_size 262144       Shape[d_model=512, heads=512]                               \n",
      "INFO:tensorflow:Variable encoder/block_003/layer_000/layer_norm/scale                 size 512          slice_size 512          Shape[d_model=512]                                          \n",
      "INFO:tensorflow:Variable encoder/block_003/layer_001/DenseReluDense/wi/kernel         size 1048576      slice_size 1048576      Shape[d_model=512, d_ff=2048]                               \n",
      "INFO:tensorflow:Variable encoder/block_003/layer_001/DenseReluDense/wo/kernel         size 1048576      slice_size 1048576      Shape[d_ff=2048, d_model=512]                               \n",
      "INFO:tensorflow:Variable encoder/block_003/layer_001/layer_norm/scale                 size 512          slice_size 512          Shape[d_model=512]                                          \n",
      "INFO:tensorflow:Variable encoder/block_004/layer_000/SelfAttention/k                  size 262144       slice_size 262144       Shape[d_model=512, heads=512]                               \n",
      "INFO:tensorflow:Variable encoder/block_004/layer_000/SelfAttention/o                  size 262144       slice_size 262144       Shape[heads=512, d_model=512]                               \n",
      "INFO:tensorflow:Variable encoder/block_004/layer_000/SelfAttention/q                  size 262144       slice_size 262144       Shape[d_model=512, heads=512]                               \n",
      "INFO:tensorflow:Variable encoder/block_004/layer_000/SelfAttention/v                  size 262144       slice_size 262144       Shape[d_model=512, heads=512]                               \n",
      "INFO:tensorflow:Variable encoder/block_004/layer_000/layer_norm/scale                 size 512          slice_size 512          Shape[d_model=512]                                          \n",
      "INFO:tensorflow:Variable encoder/block_004/layer_001/DenseReluDense/wi/kernel         size 1048576      slice_size 1048576      Shape[d_model=512, d_ff=2048]                               \n",
      "INFO:tensorflow:Variable encoder/block_004/layer_001/DenseReluDense/wo/kernel         size 1048576      slice_size 1048576      Shape[d_ff=2048, d_model=512]                               \n",
      "INFO:tensorflow:Variable encoder/block_004/layer_001/layer_norm/scale                 size 512          slice_size 512          Shape[d_model=512]                                          \n",
      "INFO:tensorflow:Variable encoder/block_005/layer_000/SelfAttention/k                  size 262144       slice_size 262144       Shape[d_model=512, heads=512]                               \n",
      "INFO:tensorflow:Variable encoder/block_005/layer_000/SelfAttention/o                  size 262144       slice_size 262144       Shape[heads=512, d_model=512]                               \n",
      "INFO:tensorflow:Variable encoder/block_005/layer_000/SelfAttention/q                  size 262144       slice_size 262144       Shape[d_model=512, heads=512]                               \n",
      "INFO:tensorflow:Variable encoder/block_005/layer_000/SelfAttention/v                  size 262144       slice_size 262144       Shape[d_model=512, heads=512]                               \n",
      "INFO:tensorflow:Variable encoder/block_005/layer_000/layer_norm/scale                 size 512          slice_size 512          Shape[d_model=512]                                          \n",
      "INFO:tensorflow:Variable encoder/block_005/layer_001/DenseReluDense/wi/kernel         size 1048576      slice_size 1048576      Shape[d_model=512, d_ff=2048]                               \n",
      "INFO:tensorflow:Variable encoder/block_005/layer_001/DenseReluDense/wo/kernel         size 1048576      slice_size 1048576      Shape[d_ff=2048, d_model=512]                               \n",
      "INFO:tensorflow:Variable encoder/block_005/layer_001/layer_norm/scale                 size 512          slice_size 512          Shape[d_model=512]                                          \n",
      "INFO:tensorflow:Variable encoder/final_layer_norm/scale                               size 512          slice_size 512          Shape[d_model=512]                                          \n",
      "INFO:tensorflow:Variable shared/embedding                                             size 16449536     slice_size 16449536     Shape[vocab=32128, d_model=512]                             \n",
      "INFO:tensorflow:Trainable Variables            count: 131     Total size: 60506624         Total slice_size: 60506624       \n",
      "INFO:tensorflow:All Variables                  count: 131     Total size: 60506624         Total slice_size: 60506624       \n",
      "INFO:tensorflow:Counters:\n",
      "einsum: 2.54e+10\n",
      "einsum_unique: 2.54e+10\n",
      "output: 2.73e+08\n",
      " output/AddOperation: 4.05e+07\n",
      " output/BinaryOpWithBroadcasting: 1.97e+06\n",
      " output/Constant: 1.57e+06\n",
      " output/EinsumOperation: 5.83e+07\n",
      " output/ImportOperation: 313\n",
      " output/MinMaxOperation: 1.18e+06\n",
      " output/OneHotOperation: 4.16e+07\n",
      " output/RangeOperation: 512\n",
      " output/ReduceOperation: 8.19e+04\n",
      " output/ReshapeOperation: 1.35e+07\n",
      " output/ScalarAddOperation: 1.58e+06\n",
      " output/ScalarMultiplyOperation: 3.68e+06\n",
      " output/ShiftOperation: 256\n",
      " output/SlicewiseOperation: 3.74e+07\n",
      " output/StopGradient: 9.44e+06\n",
      " output/Variable: 6.05e+07\n",
      " output/WhileLoopOperation: 1.57e+06\n",
      "output_unique: 2.73e+08\n",
      " output_unique/AddOperation: 4.05e+07\n",
      " output_unique/BinaryOpWithBroadcasting: 1.97e+06\n",
      " output_unique/Constant: 1.57e+06\n",
      " output_unique/EinsumOperation: 5.83e+07\n",
      " output_unique/ImportOperation: 313\n",
      " output_unique/MinMaxOperation: 1.18e+06\n",
      " output_unique/OneHotOperation: 4.16e+07\n",
      " output_unique/RangeOperation: 512\n",
      " output_unique/ReduceOperation: 8.19e+04\n",
      " output_unique/ReshapeOperation: 1.35e+07\n",
      " output_unique/ScalarAddOperation: 1.58e+06\n",
      " output_unique/ScalarMultiplyOperation: 3.68e+06\n",
      " output_unique/ShiftOperation: 256\n",
      " output_unique/SlicewiseOperation: 3.74e+07\n",
      " output_unique/StopGradient: 9.44e+06\n",
      " output_unique/Variable: 6.05e+07\n",
      " output_unique/WhileLoopOperation: 1.57e+06\n",
      "variables: 6.05e+07\n",
      " variables/trainable: 6.05e+07\n",
      "WARNING:tensorflow:From /home/husein/.local/lib/python3.6/site-packages/tensorflow_core/python/ops/array_ops.py:1475: where (from tensorflow.python.ops.array_ops) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use tf.where in 2.0, which has the same broadcast rule as np.where\n",
      "INFO:tensorflow:Done calling model_fn.\n",
      "WARNING:tensorflow:From /home/husein/.local/lib/python3.6/site-packages/tensorflow_core/python/saved_model/signature_def_utils_impl.py:201: build_tensor_info (from tensorflow.python.saved_model.utils_impl) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "This function will only be available through the v1 compatibility library as tf.compat.v1.saved_model.utils.build_tensor_info or tf.compat.v1.saved_model.build_tensor_info.\n",
      "INFO:tensorflow:Signatures INCLUDED in export for Classify: None\n",
      "INFO:tensorflow:Signatures INCLUDED in export for Regress: None\n",
      "INFO:tensorflow:Signatures INCLUDED in export for Predict: ['serving_default']\n",
      "INFO:tensorflow:Signatures INCLUDED in export for Train: None\n",
      "INFO:tensorflow:Signatures INCLUDED in export for Eval: None\n",
      "INFO:tensorflow:Restoring parameters from t5-small-spelling/model.ckpt-1100000\n",
      "INFO:tensorflow:Assets added to graph.\n",
      "INFO:tensorflow:No assets to write.\n",
      "INFO:tensorflow:SavedModel written to: output/temp-b'1631417911'/saved_model.pb\n"
     ]
    }
   ],
   "source": [
    "from mesh_tensorflow.transformer import dataset as transformer_dataset\n",
    "\n",
    "def serving_input_fn():\n",
    "    inputs = tf.placeholder(\n",
    "            dtype=tf.string,\n",
    "            shape=[None],\n",
    "            name=\"inputs\")\n",
    "\n",
    "    batch_size = tf.shape(inputs)[0]\n",
    "    padded_inputs = tf.pad(inputs, [(0, tf.mod(-tf.size(inputs), batch_size))])\n",
    "    dataset = tf.data.Dataset.from_tensor_slices(padded_inputs)\n",
    "    dataset = dataset.map(lambda x: {\"inputs\": x})\n",
    "    dataset = transformer_dataset.encode_all_features(dataset, vocabulary)\n",
    "    dataset = transformer_dataset.pack_or_pad(\n",
    "        dataset=dataset,\n",
    "        length=model._sequence_length,\n",
    "        pack=False,\n",
    "        feature_keys=[\"inputs\"]\n",
    "    )\n",
    "    dataset = dataset.batch(tf.cast(batch_size, tf.int64))\n",
    "    features = tf.data.experimental.get_single_element(dataset)\n",
    "    return tf.estimator.export.ServingInputReceiver(\n",
    "        features=features, receiver_tensors=inputs)\n",
    "\n",
    "out = estimator.export_saved_model('output', serving_input_fn, checkpoint_path=checkpoint_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From <ipython-input-16-5b89b6a20c22>:7: load (from tensorflow.python.saved_model.loader_impl) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "This function will only be available through the v1 compatibility library as tf.compat.v1.saved_model.loader.load or tf.compat.v1.saved_model.load. There will be a new function for importing SavedModels in Tensorflow 2.0.\n",
      "INFO:tensorflow:Restoring parameters from output/1631417911/variables/variables\n"
     ]
    }
   ],
   "source": [
    "config = tf.ConfigProto()\n",
    "config.allow_soft_placement = True\n",
    "sess = tf.Session(config = config)\n",
    "meta_graph_def = tf.saved_model.loader.load(\n",
    "        sess,\n",
    "        [tf.saved_model.tag_constants.SERVING],\n",
    "        out)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'small-spelling-correction/model.ckpt'"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "saver = tf.train.Saver(tf.trainable_variables())\n",
    "saver.save(sess, 'small-spelling-correction/model.ckpt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "strings = [\n",
    "    n.name\n",
    "    for n in tf.get_default_graph().as_graph_def().node\n",
    "    if ('encoder' in n.op\n",
    "    or 'decoder' in n.name\n",
    "    or 'shared' in n.name\n",
    "    or 'inputs' in n.name\n",
    "    or 'output' in n.name\n",
    "    or 'SentenceTokenizer' in n.name\n",
    "    or 'self/Softmax' in n.name)\n",
    "    and 'adam' not in n.name\n",
    "    and 'Assign' not in n.name\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "def freeze_graph(model_dir, output_node_names):\n",
    "\n",
    "    if not tf.gfile.Exists(model_dir):\n",
    "        raise AssertionError(\n",
    "            \"Export directory doesn't exists. Please specify an export \"\n",
    "            'directory: %s' % model_dir\n",
    "        )\n",
    "\n",
    "    checkpoint = tf.train.get_checkpoint_state(model_dir)\n",
    "    input_checkpoint = checkpoint.model_checkpoint_path\n",
    "\n",
    "    absolute_model_dir = '/'.join(input_checkpoint.split('/')[:-1])\n",
    "    output_graph = absolute_model_dir + '/frozen_model.pb'\n",
    "    clear_devices = True\n",
    "    with tf.Session(graph = tf.Graph()) as sess:\n",
    "        saver = tf.train.import_meta_graph(\n",
    "            input_checkpoint + '.meta', clear_devices = clear_devices\n",
    "        )\n",
    "        saver.restore(sess, input_checkpoint)\n",
    "        output_graph_def = tf.graph_util.convert_variables_to_constants(\n",
    "            sess,\n",
    "            tf.get_default_graph().as_graph_def(),\n",
    "            output_node_names,\n",
    "        )\n",
    "        with tf.gfile.GFile(output_graph, 'wb') as f:\n",
    "            f.write(output_graph_def.SerializeToString())\n",
    "        print('%d ops in the final graph.' % len(output_graph_def.node))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Restoring parameters from small-spelling-correction/model.ckpt\n",
      "WARNING:tensorflow:From <ipython-input-19-504c79665720>:23: convert_variables_to_constants (from tensorflow.python.framework.graph_util_impl) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use `tf.compat.v1.graph_util.convert_variables_to_constants`\n",
      "WARNING:tensorflow:From /home/husein/.local/lib/python3.6/site-packages/tensorflow_core/python/framework/graph_util_impl.py:277: extract_sub_graph (from tensorflow.python.framework.graph_util_impl) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use `tf.compat.v1.graph_util.extract_sub_graph`\n",
      "INFO:tensorflow:Froze 212 variables.\n",
      "INFO:tensorflow:Converted 212 variables to const ops.\n",
      "7064 ops in the final graph.\n"
     ]
    }
   ],
   "source": [
    "freeze_graph('small-spelling-correction', strings)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "import struct\n",
    "\n",
    "unknown = b'\\xff\\xff\\xff\\xff'\n",
    "\n",
    "def load_graph(frozen_graph_filename):\n",
    "    with tf.gfile.GFile(frozen_graph_filename, 'rb') as f:\n",
    "        graph_def = tf.GraphDef()\n",
    "        graph_def.ParseFromString(f.read())\n",
    "        \n",
    "    for node in graph_def.node:\n",
    "        \n",
    "        if node.op == 'RefSwitch':\n",
    "          node.op = 'Switch'\n",
    "          for index in xrange(len(node.input)):\n",
    "            if 'moving_' in node.input[index]:\n",
    "              node.input[index] = node.input[index] + '/read'\n",
    "        elif node.op == 'AssignSub':\n",
    "          node.op = 'Sub'\n",
    "          if 'use_locking' in node.attr: del node.attr['use_locking']\n",
    "        elif node.op == 'AssignAdd':\n",
    "          node.op = 'Add'\n",
    "          if 'use_locking' in node.attr: del node.attr['use_locking']\n",
    "        elif node.op == 'Assign':\n",
    "          node.op = 'Identity'\n",
    "          if 'use_locking' in node.attr: del node.attr['use_locking']\n",
    "          if 'validate_shape' in node.attr: del node.attr['validate_shape']\n",
    "          if len(node.input) == 2:\n",
    "            node.input[0] = node.input[1]\n",
    "            del node.input[1]\n",
    "            \n",
    "        if 'Reshape/shape' in node.name or 'Reshape_1/shape' in node.name:\n",
    "            b = node.attr['value'].tensor.tensor_content\n",
    "            arr_int = [int.from_bytes(b[i:i + 4], 'little') for i in range(0, len(b), 4)]\n",
    "            if len(arr_int):\n",
    "                arr_byte = [unknown] + [struct.pack('<i', i) for i in arr_int[1:]]\n",
    "                arr_byte = b''.join(arr_byte)\n",
    "                node.attr['value'].tensor.tensor_content = arr_byte\n",
    "            \n",
    "            if len(node.attr['value'].tensor.int_val):\n",
    "                node.attr['value'].tensor.int_val[0] = -1\n",
    "    \n",
    "    with tf.Graph().as_default() as graph:\n",
    "        tf.import_graph_def(graph_def)\n",
    "    return graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "g = load_graph('small-spelling-correction/frozen_model.pb')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(<tf.Tensor 'import/inputs:0' shape=(?,) dtype=string>,\n",
       " <tf.Tensor 'import/SelectV2_3:0' shape=(?, 256) dtype=int32>)"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "i = g.get_tensor_by_name('import/inputs:0')\n",
    "o = g.get_tensor_by_name('import/SelectV2_3:0')\n",
    "i, o"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_sess = tf.Session(graph = g)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import sentencepiece as spm\n",
    "sp_model = spm.SentencePieceProcessor()\n",
    "sp_model.Load(DEFAULT_SPM_PATH)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "string1 = 'krajaan patut bagi pencen awal skt kpd warga emas supaya emosi'\n",
    "string2 = 'Husein ska mkn aym dkat kampng Jawa'\n",
    "string3 = 'Melayu malas ni narration dia sama je macam men are trash. True to some, false to some.'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['betulkan ejaan: krajaan patut bagi pencen awal skt kpd warga emas supaya emosi',\n",
       " 'betulkan ejaan: Husein ska mkn aym dkat kampng Jawa',\n",
       " 'betulkan ejaan: Melayu malas ni narration dia sama je macam men are trash. True to some, false to some.']"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "strings = [string1, string2, string3]\n",
    "[f'betulkan ejaan: {s}' for s in strings]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 14.1 s, sys: 4.79 s, total: 18.9 s\n",
      "Wall time: 9.02 s\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(3, 256)"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "o_ = test_sess.run(o, feed_dict = {i: [f'betulkan ejaan: {s}' for s in strings]})\n",
    "o_.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 krajaan patut bagi pencen awal skt kpd warga emas supaya emosi\n",
      "1 Husein sukah makan ayam dekat kampung Jawa\n",
      "2 Melayu malas ni narration dia sama je macam men are trash. True to some, false to some.\n"
     ]
    }
   ],
   "source": [
    "for k in range(len(o_)):\n",
    "    print(k, sp_model.DecodeIds(o_[k].tolist()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tensorflow.tools.graph_transforms import TransformGraph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "transforms = ['add_default_attributes',\n",
    "             'remove_nodes(op=Identity, op=CheckNumerics)',\n",
    "             'fold_batch_norms',\n",
    "             'fold_old_batch_norms',\n",
    "             'quantize_weights(minimum_size=1536000)',\n",
    "             #'quantize_weights(fallback_min=-10240, fallback_max=10240)',\n",
    "             'strip_unused_nodes',\n",
    "             'sort_by_execution_order']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From <ipython-input-33-3c601715f80e>:3: FastGFile.__init__ (from tensorflow.python.platform.gfile) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use tf.gfile.GFile.\n"
     ]
    }
   ],
   "source": [
    "pb = 'small-spelling-correction/frozen_model.pb'\n",
    "input_graph_def = tf.GraphDef()\n",
    "with tf.gfile.FastGFile(pb, 'rb') as f:\n",
    "    input_graph_def.ParseFromString(f.read())\n",
    "    \n",
    "transformed_graph_def = TransformGraph(input_graph_def, \n",
    "       ['inputs'],\n",
    "       ['SelectV2_3'], transforms)\n",
    "\n",
    "with tf.gfile.GFile(f'{pb}.quantized', 'wb') as f:\n",
    "    f.write(transformed_graph_def.SerializeToString())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(<tf.Tensor 'import/inputs:0' shape=(?,) dtype=string>,\n",
       " <tf.Tensor 'import/SelectV2_3:0' shape=(?, 256) dtype=int32>)"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "g = load_graph('small-spelling-correction/frozen_model.pb.quantized')\n",
    "i = g.get_tensor_by_name('import/inputs:0')\n",
    "o = g.get_tensor_by_name('import/SelectV2_3:0')\n",
    "i, o"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_sess = tf.InteractiveSession(graph = g)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<b2sdk.file_version.FileVersionInfo at 0x7f78c80c2550>"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "file = 'small-spelling-correction/frozen_model.pb.quantized'\n",
    "outPutname = 'spelling-correction/small-t5-quantized/model.pb'\n",
    "b2_bucket.upload_local_file(\n",
    "    local_file=file,\n",
    "    file_name=outPutname,\n",
    "    file_infos=file_info,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<b2sdk.file_version.FileVersionInfo at 0x7f78887dc048>"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "file = 'small-spelling-correction/frozen_model.pb'\n",
    "outPutname = 'spelling-correction/small-t5/model.pb'\n",
    "b2_bucket.upload_local_file(\n",
    "    local_file=file,\n",
    "    file_name=outPutname,\n",
    "    file_infos=file_info,\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
